#!/usr/bin/env python3
#-*- coding:utf8 -*-
# Power by 2020-06-13 00:48:57

import os
import random
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import numpy as np
from PIL import Image
import gzip
import json

class MNIST(fluid.dygraph.Layer):

    """Docstring for MNIST. """

    def __init__(self):
        """TODO: to be defined. """
        super(MNIST,self).__init__()
        self.conv1=Conv2D(num_channels=1,num_filters=20,filter_size=5,stride=1,padding=2,act='relu')
        self.pool1=Pool2D(pool_size=2,pool_stride=2,pool_type='max')
        self.conv2=Conv2D(num_channels=20,num_filters=20,filter_size=5,stride=1,padding=2,act='relu')
        self.pool2=Pool2D(pool_size=2,pool_stride=2,pool_type='max')
        self.fc=Linear(input_dim=980,output_dim=10,act='softmax')
    def forward(self, inputs):
        """TODO: Docstring for forward.

        :inputs: TODO
        :returns: TODO

        """
        x=self.conv1(inputs)
        x=self.pool1(x)
        x=self.conv2(x)
        x=self.pool2(x)
        x=fluid.layers.reshape(x,[x.shape[0],980])
        x=self.fc(x)
        return x
        
def load_image(img_path):
    """TODO: Docstring for load_image.

    :img_path: TODO
    :returns: TODO

    """
    im=Image.open(img_path).convert('L')
    im.show()
    im=im.resize((28,28),Image.ANTIALIAS)
    im=np.array(im).reshape(1,1,28,28).astype(np.float32)
    im=1.0 - im/255.
    return im
with fluid.dygraph.guard():
    model=MNIST()
    img_path='./example_0.jpg'
    model_dict,_=fluid.load_dygraph('./mnist_cross_entropy.pdparams')
    model.load_dict(model_dict)
    model.eval()
    tensor_img=load_image(img_path)
    results=model(fluid.dygraph.to_variable(tensor_img))
    lab=np.argsort(results.numpy())
    print("the num is:{}".format(lab[0][-1]))
